import torch
import torch.onnx
import onnx
import onnxruntime
import numpy as np
from network_files import MaskRCNN
from backbone import resnet50_fpn_backbone, resnet18_fpn_backbone
import warnings
warnings.filterwarnings("ignore")

device = torch.device("cpu")

def create_model(num_classes, box_thresh=0.5):
    # backbone = resnet50_fpn_backbone()
    backbone = resnet18_fpn_backbone()
    model = MaskRCNN(backbone,
                     num_classes=num_classes,
                     rpn_score_thresh=box_thresh,
                     box_score_thresh=box_thresh)

    return model

def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()


def main(save_path, weights_path):
    assert isinstance(save_path, str), "lack of save_path parameter..."
    print("Creating model......")
    # create model
    model = create_model(num_classes=3+ 1)

    # assert os.path.exists(weights_path), "{} file dose not exist.".format(weights_path)
    weights_dict = torch.load(weights_path, map_location='cpu')
    weights_dict = weights_dict["model"] if "model" in weights_dict else weights_dict
    model.load_state_dict(weights_dict)


    model.eval()
    # input to the model
    # [batch, channel, height, width]
    x = torch.rand(1, 3, 448, 1216, requires_grad=False)
    # torch_out = model(x)

    # export the model
    torch.onnx.export(model,                       # model being run
                      x,                           # model input (or a tuple for multiple inputs)
                      save_path,                   # where to save the model (can be a file or file-like object)
                      export_params=True,          # store the trained parameter weights inside the model file
                      opset_version=16,            # the ONNX version to export the model to
                      do_constant_folding=True,    # whether to execute constant folding for optimization
                      input_names=["input"],       # the model's input names
                      output_names=["output"],     # the model's output names
                      dynamic_axes={"input": {0: "batch_size"},  # variable length axes
                                    "output": {0: "batch_size"}})

    # check onnx model
    onnx_model = onnx.load(save_path)
    onnx.checker.check_model(onnx_model)
    print("-"*30)
    print("|  onnx conver sucess!!!             |")
    print("|  onnx saved to: {}  |".format(onnx_file_name))
    print("-" * 30)
if __name__ == '__main__':
    onnx_file_name = "onnx/maskrcnn_r18.onnx"
    weights_path = "./multi_train_r18/model_35.pth"
    main(save_path=onnx_file_name, weights_path=weights_path)
